// syscall_monitor.c
#include <bpf/libbpf.h>
#include <bpf/bpf.h>
#include "syscall_monitor.skel.h"

#include <sys/resource.h>
#include <unistd.h>
#include <signal.h>
#include <getopt.h>
#include "syscall_monitor.h"

// libbpf日志回调函数
static int libbpf_print_fn(enum libbpf_print_level level,
                           const char* format,
                           va_list args) {
    return vfprintf(stderr, format, args);
}

// 信号处理函数，用于优雅退出
static volatile bool exiting = false;
static void sig_handler(int sig) {
    exiting = true;
}

// 打印进程信息
static void print_process_info(struct process_info *info) {
    printf("进程信息 - PID: %u, 名称: %s, 系统调用次数: %llu, 启动时间: %llu\n",
           info->pid, info->comm, info->syscall_count, info->start_time);
}

// 打印系统调用事件
static void print_syscall_event(struct syscall_event *event) {
    printf("系统调用事件 - PID: %u, 名称: %s, 系统调用号: %u, 时间戳: %llu\n",
           event->pid, event->comm, event->syscall_nr, event->timestamp);
}

// 打印进程生命周期事件
static void print_process_lifecycle_event(struct process_lifecycle_event *event) {
    const char *event_types[] = {"创建", "退出"};
    printf("进程生命周期事件 - PID: %u, 名称: %s, 事件类型: %s, 时间戳: %llu\n",
           event->pid, event->comm, 
           event->event_type < 2 ? event_types[event->event_type] : "未知",
           event->timestamp);
}

// 打印CPU调度事件
static void print_cpu_sched_event(struct cpu_sched_info *event) {
    printf("CPU调度事件 - PID: %u, 名称: %s, CPU: %u, 时间戳: %llu\n",
           event->pid, event->comm, event->cpu, event->timestamp);
}

// 打印内存事件
static void print_memory_event(struct memory_event *event) {
    const char *event_types[] = {"kmalloc", "mmap", "slab_alloc"};
    printf("内存事件 - PID: %u, 名称: %s, 事件类型: %s, 大小: %llu, 时间戳: %llu\n",
           event->pid, event->comm,
           event->event_type < 3 ? event_types[event->event_type] : "未知",
           event->size, event->timestamp);
}

// 打印文件IO事件
static void print_file_io_event(struct file_io_event *event) {
    const char *operations[] = {"读取", "写入"};
    printf("文件IO事件 - PID: %u, 名称: %s, 操作: %s, 字节数: %llu, 时间戳: %llu, 文件名: %s\n",
           event->pid, event->comm,
           event->operation < 2 ? operations[event->operation] : "未知",
           event->bytes, event->timestamp, event->filename);
}

// 打印网络连接事件
static void print_network_event(struct network_connection *event) {
    const char *event_types[] = {"连接", "关闭", "状态变更"};
    const char *protocols[] = {"TCP", "UDP"};
    printf("网络事件 - PID: %u, 名称: %s, 协议: %s, 事件类型: %s, 时间戳: %llu\n",
           event->pid, event->comm,
           event->protocol <= 1 ? protocols[event->protocol] : "未知",
           event->event_type < 3 ? event_types[event->event_type] : "未知",
           event->timestamp);
}

// 打印锁竞争事件
static void print_lock_contention_event(struct lock_contention_event *event) {
    const char *lock_types[] = {"mutex", "spinlock"};
    printf("锁竞争事件 - PID: %u, 名称: %s, 锁类型: %s, 等待时间: %llu, 时间戳: %llu\n",
           event->pid, event->comm,
           event->lock_type < 2 ? lock_types[event->lock_type] : "未知",
           event->wait_time, event->timestamp);
}

// 打印中断事件
static void print_interrupt_event(struct interrupt_event *event) {
    const char *event_types[] = {"硬件", "软中断"};
    printf("中断事件 - 中断号: %u, CPU: %u, 事件类型: %s, 处理耗时: %llu, 时间戳: %llu\n",
           event->irq_num, event->cpu,
           event->event_type < 2 ? event_types[event->event_type] : "未知",
           event->duration, event->timestamp);
}

// 打印页错误事件
static void print_page_fault_event(struct page_fault_event *event) {
    printf("页错误事件 - PID: %u, 名称: %s, 地址: %llu, 错误类型: %u, 时间戳: %llu\n",
           event->pid, event->comm, event->address, event->fault_type, event->timestamp);
}

// 打印安全事件
static void print_security_event(struct security_event *event) {
    const char *event_types[] = {"文件访问", "权限变更", "模块加载"};
    printf("安全事件 - PID: %u, 名称: %s, 事件类型: %s, 时间戳: %llu, 路径: %s\n",
           event->pid, event->comm,
           event->event_type < 3 ? event_types[event->event_type] : "未知",
           event->timestamp, event->path);
}

// 打印容器事件
static void print_container_event(struct container_event *event) {
    const char *event_types[] = {"创建", "销毁", "资源"};
    printf("容器事件 - PID: %u, 名称: %s, 事件类型: %s, cgroup ID: %llu, 时间戳: %llu\n",
           event->pid, event->comm,
           event->event_type < 3 ? event_types[event->event_type] : "未知",
           event->cgroup_id, event->timestamp);
}

#ifdef OUTPUT_MODE_PERF_EVENT
// Perf event 回调函数
static void handle_syscall_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct syscall_event *event = (struct syscall_event *)data;
    print_syscall_event(event);
}

static void handle_process_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct process_lifecycle_event *event = (struct process_lifecycle_event *)data;
    print_process_lifecycle_event(event);
}

static void handle_cpu_sched_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct cpu_sched_info *event = (struct cpu_sched_info *)data;
    print_cpu_sched_event(event);
}

static void handle_memory_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct memory_event *event = (struct memory_event *)data;
    print_memory_event(event);
}

static void handle_file_io_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct file_io_event *event = (struct file_io_event *)data;
    print_file_io_event(event);
}

static void handle_network_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct network_connection *event = (struct network_connection *)data;
    print_network_event(event);
}

static void handle_lock_contention_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct lock_contention_event *event = (struct lock_contention_event *)data;
    print_lock_contention_event(event);
}

static void handle_interrupt_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct interrupt_event *event = (struct interrupt_event *)data;
    print_interrupt_event(event);
}

static void handle_page_fault_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct page_fault_event *event = (struct page_fault_event *)data;
    print_page_fault_event(event);
}

static void handle_security_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct security_event *event = (struct security_event *)data;
    print_security_event(event);
}

static void handle_container_event(void *ctx, int cpu, void *data, __u32 data_sz) {
    struct container_event *event = (struct container_event *)data;
    print_container_event(event);
}
#endif

#ifdef OUTPUT_MODE_MAP_POLL
// 轮询并打印映射中的事件数据
static void poll_and_print_events(struct syscall_monitor_bpf *skel) {
    printf("\n=== 进程信息 ===\n");
    unsigned int key, next_key;
    struct process_info info;

    key = 0;
    int count = 0;
    while (bpf_map__get_next_key(
           skel->maps.process_map,  // 直接使用 struct bpf_map *
           count == 0 ? NULL : &key, 
           &next_key,
           sizeof(next_key)) == 0) {
        key = next_key;
        if (bpf_map_lookup_elem(bpf_map__fd(skel->maps.process_map), &key, &info) == 0) {
            print_process_info(&info);
        }
        count++;
        if (count >= 10) break; // 限制输出数量
    }

    if (count == 0) {
        printf("暂无进程信息\n");
    }
}
#endif

void print_usage(const char *program_name) {
    printf("用法: %s [选项]\n", program_name);
    printf("选项:\n");
    printf("  -h, --help              显示此帮助信息\n");
    printf("  -a, --all-processes     监控所有进程 (默认行为)\n");
    printf("  -p, --pid PID           只监控指定PID的进程\n");
    printf("\n");
}

int main(int argc, char** argv) {
    struct syscall_monitor_bpf* skel;
    int err;
    int target_pid = 0; // 0表示监控所有进程

    // 解析命令行参数
    static struct option long_options[] = {
        {"help",          no_argument,       0, 'h'},
        {"all-processes", no_argument,       0, 'a'},
        {"pid",           required_argument, 0, 'p'},
        {0, 0, 0, 0}
    };

    int opt;
    while ((opt = getopt_long(argc, argv, "hap:", long_options, NULL)) != -1) {
        switch (opt) {
        case 'h':
            print_usage(argv[0]);
            return 0;
        case 'a':
            target_pid = 0; // 监控所有进程
            break;
        case 'p':
            target_pid = atoi(optarg);
            if (target_pid <= 0) {
                fprintf(stderr, "无效的PID: %s\n", optarg);
                return 1;
            }
            break;
        default:
            print_usage(argv[0]);
            return 1;
        }
    }

    // 设置libbpf的错误和调试信息回调函数
    libbpf_set_print(libbpf_print_fn);

    // 设置信号处理
    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);

    // 设置资源限制
    struct rlimit rl = {.rlim_cur = RLIM_INFINITY, .rlim_max = RLIM_INFINITY};
    if (setrlimit(RLIMIT_MEMLOCK, &rl)) {
        fprintf(stderr, "无法设置 RLIMIT_MEMLOCK 限制: %s\n", strerror(errno));
        return 1;
    }

    // 打开eBPF应用程序
    skel = syscall_monitor_bpf__open();
    if (!skel) {
        fprintf(stderr, "无法打开eBPF程序\n");
        return 1;
    }

    // 设置要监控的PID（0表示监控所有进程）
    skel->bss->my_pid = target_pid;

    // 加载并验证eBPF程序
    err = syscall_monitor_bpf__load(skel);
    if (err) {
        fprintf(stderr, "加载和验证eBPF程序失败: %d\n", err);
        goto cleanup;
    }

    // 挂载到tracepoint
    err = syscall_monitor_bpf__attach(skel);
    if (err) {
        fprintf(stderr, "附加eBPF程序失败: %d\n", err);
        goto cleanup;
    }

#ifdef OUTPUT_MODE_PERF_EVENT
    // 设置perf event回调
    struct perf_buffer *pb = NULL;
    struct perf_buffer_opts pb_opts = {};
    
    // 创建perf buffer并注册回调函数
    pb = perf_buffer__new(bpf_map__fd(skel->maps.syscall_events), 8, 
                          handle_syscall_event, NULL, NULL, &pb_opts);
    if (!pb) {
        err = -1;
        fprintf(stderr, "创建syscall_events perf buffer失败\n");
        goto cleanup;
    }
    
    // 为其他事件类型注册回调
    perf_buffer__new(bpf_map__fd(skel->maps.process_events), 8,
                     handle_process_event, NULL, NULL, &pb_opts);
    perf_buffer__new(bpf_map__fd(skel->maps.cpu_sched_events), 8,
                     handle_cpu_sched_event, NULL, NULL, &pb_opts);
    perf_buffer__new(bpf_map__fd(skel->maps.memory_events), 8,
                     handle_memory_event, NULL, NULL, &pb_opts);
    perf_buffer__new(bpf_map__fd(skel->maps.file_io_events), 8,
                     handle_file_io_event, NULL, NULL, &pb_opts);
    perf_buffer__new(bpf_map__fd(skel->maps.network_events), 8,
                     handle_network_event, NULL, NULL, &pb_opts);
    perf_buffer__new(bpf_map__fd(skel->maps.lock_contention_events), 8,
                     handle_lock_contention_event, NULL, NULL, &pb_opts);
    perf_buffer__new(bpf_map__fd(skel->maps.interrupt_events), 8,
                     handle_interrupt_event, NULL, NULL, &pb_opts);
    perf_buffer__new(bpf_map__fd(skel->maps.page_fault_events), 8,
                     handle_page_fault_event, NULL, NULL, &pb_opts);
    perf_buffer__new(bpf_map__fd(skel->maps.security_events), 8,
                     handle_security_event, NULL, NULL, &pb_opts);
    perf_buffer__new(bpf_map__fd(skel->maps.container_events), 8,
                     handle_container_event, NULL, NULL, &pb_opts);

    if (target_pid == 0) {
        printf("成功启动eBPF监控程序! 将监控所有进程\n");
    } else {
        printf("成功启动eBPF监控程序! 只监控PID为 %d 的进程\n", target_pid);
    }
    printf("程序将实时输出收集到的事件信息...\n");
    printf("按 Ctrl+C 停止程序\n");

    // 主循环 - 保持程序运行并处理事件
    while (!exiting) {
        // 轮询perf buffer处理事件
        err = perf_buffer__poll(pb, 100);
        if (err == -EINTR) {
            err = 0;
            break;
        }
        if (err) {
            fprintf(stderr, "perf_buffer__poll 错误: %d\n", err);
            break;
        }
    }

#elif defined(OUTPUT_MODE_MAP_POLL)
    if (target_pid == 0) {
        printf("成功启动eBPF监控程序! 将监控所有进程\n");
    } else {
        printf("成功启动eBPF监控程序! 只监控PID为 %d 的进程\n", target_pid);
    }
    printf("程序将每5秒输出一次收集到的事件信息...\n");
    printf("按 Ctrl+C 停止程序\n");

    // 主循环 - 保持程序运行并定期输出事件信息
    while (!exiting) {
        // 每5秒输出一次事件信息
        sleep(5);
        
        // 轮询并打印事件
        poll_and_print_events(skel);
        
        // 触发一些系统调用来产生监控数据
        FILE* f = fopen("/proc/self/comm", "r");
        if (f) {
            char buf[16];
            fread(buf, 1, sizeof(buf), f);
            fclose(f);
        }
    }
#endif

    printf("正在退出...\n");

cleanup:
#ifdef OUTPUT_MODE_PERF_EVENT
    // 清理资源
    if (pb) {
        perf_buffer__free(pb);
    }
#endif
    syscall_monitor_bpf__destroy(skel);
    return -err;
}